//
//  CollectionFlowLayout.swift
//  Risac
//
//  Created by remy on 2022/12/23.
//  Copyright © 2022 com.Risac. All rights reserved.
//

import UIKit

open class CollectionFlowLayout: UICollectionViewLayout {
    public weak var model: CollectionModel?
    public let direction: UICollectionView.ScrollDirection
    private var cellAttrs: [IndexPath: UICollectionViewLayoutAttributes] = [:]
    private var sectionAttrs: [Int: UICollectionViewLayoutAttributes] = [:]
    private var contentSize: CGSize = .zero
    override open var collectionViewContentSize: CGSize {
        return self.contentSize
    }
    
    public init(model: CollectionModel,
                direction: UICollectionView.ScrollDirection = .vertical) {
        self.model = model
        self.direction = direction
        super.init()
    }
    
    required public init?(coder: NSCoder) {
        fatalError("init(coder:) has not been implemented")
    }
    
    override open func prepare() {
        super.prepare()
        prepareAttributes()
    }
    
    private func prepareAttributes() {
        guard let model = model else { return }
        guard let collectionView = self.collectionView else { return }
        sectionAttrs.removeAll()
        cellAttrs.removeAll()
        let viewSize: CGSize = collectionView.size
        var sectionStartPoint: CGPoint = .zero
        var offset: CGPoint = .zero
        var tracks: [CGFloat] = []
        if direction == .vertical {
            model.sectionItems.forEachEnumerated { (sIndex, section) in
                let indexPath = IndexPath(row: 0, section: sIndex)
                let sectionMargin = section.sectionMargin
                let sectionInset = section.sectionInset
                let lineSpacing = section.lineSpacing
                let itemSpacing = section.itemSpacing
                let flowCount = max(section.flowCount, 1)
                let attr = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: UICollectionView.elementKindSectionStage, with: indexPath)
                attr.zIndex = -1
                sectionAttrs[sIndex] = attr
                let trackSize = (viewSize.width - sectionInset.left - sectionInset.right - sectionMargin.left - sectionMargin.right - itemSpacing * CGFloat(flowCount - 1)) / CGFloat(flowCount).rounded()
                sectionStartPoint = CGPoint(x: sectionMargin.left, y: (sIndex > 0 ? sectionStartPoint.y + section.sectionSpacing : 0.0) + sectionMargin.top)
                tracks = Array(repeating: sectionStartPoint.y + sectionInset.top, count: flowCount)
                section.cellItems.forEachEnumerated { (cIndex, cell) in
                    let indexPath = IndexPath(row: cIndex, section: sIndex)
                    var cellSize = cell.cellSize
                    cellSize.height = (cellSize.height / cellSize.width * trackSize).rounded()
                    cellSize.width = trackSize
                    var insertIndex: Int = 0
                    tracks.forEachEnumerated {
                        insertIndex = $1 < tracks[insertIndex] ? $0 : insertIndex
                    }
                    offset.x = sectionMargin.left + sectionInset.left + (trackSize + itemSpacing) * CGFloat(insertIndex)
                    offset.y = tracks[insertIndex] + (cIndex < flowCount ? 0.0 : lineSpacing)
                    let attr = UICollectionViewLayoutAttributes(forCellWith: indexPath)
                    attr.frame = CGRect(origin: offset, size: cellSize)
                    cellAttrs[indexPath] = attr
                    tracks[insertIndex] = attr.frame.maxY
                }
                attr.frame = CGRect(origin: sectionStartPoint, size: CGSize(width: viewSize.width - sectionMargin.left - sectionMargin.right, height: tracks.reduce(0, { max($0, $1) }) - sectionStartPoint.y + sectionInset.bottom))
                sectionStartPoint.y = attr.frame.maxY + sectionMargin.bottom
                contentSize = CGSize(width: viewSize.width, height: sectionStartPoint.y)
            }
        } else {
            model.sectionItems.forEachEnumerated { (sIndex, section) in
                let indexPath = IndexPath(row: 0, section: sIndex)
                let sectionMargin = section.sectionMargin
                let sectionInset = section.sectionInset
                let lineSpacing = section.lineSpacing
                let itemSpacing = section.itemSpacing
                let flowCount = max(section.flowCount, 1)
                let attr = UICollectionViewLayoutAttributes(forSupplementaryViewOfKind: UICollectionView.elementKindSectionStage, with: indexPath)
                attr.zIndex = -1
                sectionAttrs[sIndex] = attr
                let trackSize = (viewSize.height - sectionInset.top - sectionInset.bottom - sectionMargin.top - sectionMargin.bottom - lineSpacing * CGFloat(flowCount - 1)) / CGFloat(flowCount).rounded()
                sectionStartPoint = CGPoint(x: (sIndex > 0 ? sectionStartPoint.x + section.sectionSpacing : 0.0) + sectionMargin.left, y: sectionMargin.top)
                tracks = Array(repeating: sectionStartPoint.x + sectionInset.left, count: flowCount)
                section.cellItems.forEachEnumerated { (cIndex, cell) in
                    let indexPath = IndexPath(row: cIndex, section: sIndex)
                    var cellSize = cell.cellSize
                    cellSize.width = (cellSize.width / cellSize.height * trackSize).rounded()
                    cellSize.height = trackSize
                    var insertIndex: Int = 0
                    tracks.forEachEnumerated {
                        insertIndex = $1 < tracks[insertIndex] ? $0 : insertIndex
                    }
                    offset.x = tracks[insertIndex] + (cIndex < flowCount ? 0.0 : itemSpacing)
                    offset.y = sectionMargin.top + sectionInset.top + (trackSize + lineSpacing) * CGFloat(insertIndex)
                    let attr = UICollectionViewLayoutAttributes(forCellWith: indexPath)
                    attr.frame = CGRect(origin: offset, size: cellSize)
                    cellAttrs[indexPath] = attr
                    tracks[insertIndex] = attr.frame.maxX
                }
                attr.frame = CGRect(origin: sectionStartPoint, size: CGSize(width: tracks.reduce(0, { max($0, $1) }) - sectionStartPoint.x + sectionInset.right, height: viewSize.height - sectionMargin.top - sectionMargin.bottom))
                sectionStartPoint.x = attr.frame.maxX + sectionMargin.right
                contentSize = CGSize(width: sectionStartPoint.x, height: viewSize.height)
            }
        }
    }
    
    override open func layoutAttributesForItem(at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        return cellAttrs[indexPath]
    }
    
    override open func layoutAttributesForElements(in rect: CGRect) -> [UICollectionViewLayoutAttributes]? {
        return cellAttrs.values.filter { $0.frame.intersects(rect) } + sectionAttrs.values.filter { $0.frame.intersects(rect) }
    }
    
    override open func layoutAttributesForSupplementaryView(ofKind elementKind: String, at indexPath: IndexPath) -> UICollectionViewLayoutAttributes? {
        if elementKind == UICollectionView.elementKindSectionStage {
            return sectionAttrs[indexPath.section]
        }
        return super.layoutAttributesForSupplementaryView(ofKind: elementKind, at: indexPath)
    }
}
